package main

import (
	"sync"
)

// Tasks 结构体
type Tasks[T any] struct {
	numGoroutines int
	taskChannel   chan *T
	wg            sync.WaitGroup
}

// 新建 Tasks 实例
func NewTasks[T any](numGoroutines int) *Tasks[T] {
	return &Tasks[T]{
		numGoroutines: numGoroutines,
		taskChannel:   make(chan *T, numGoroutines),
	}
}

// 启动协程
func (t *Tasks[T]) Start(handler func(t *T)) {
	t.wg.Add(t.numGoroutines)

	for i := 0; i < t.numGoroutines; i++ {
		go func() {
			defer t.wg.Done()

			for task := range t.taskChannel {
				if task == nil {
					break
				}
				// 处理任务
				handler(task)
			}
		}()
	}
}

// 发送任务
func (t *Tasks[T]) SendTask(task *T) {
	t.taskChannel <- task
}

// 关闭通道，通知协程结束
func (t *Tasks[T]) Close() {
	for i := 0; i < t.numGoroutines; i++ {
		t.SendTask(nil)
	}
}

// 等待所有协程完成
func (t *Tasks[T]) Wait() {
	t.wg.Wait()
}
